package com.xes.hardwritingrecognition.pytorch;

import android.content.Context;

import com.xes.hardwritingrecognition.Classifier;

import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * @author wl
 * @since 2020/09/25 17:27
 */
public class PytorchClassifier implements Classifier {
    public static final String TAG = "PytorchClassifier";
    private boolean isLog = false;
    Module module;
    int inputsize;

    private PytorchClassifier() {
    }

    public static PytorchClassifier create(Context context, InputStream inputStream, int inputsize) throws IOException {
        PytorchClassifier p = new PytorchClassifier();
        p.module = Module.load(saveStreamAsFile(context, inputStream, new File(context.getFilesDir(), "pytorch_def")));
        p.inputsize = inputsize;
        return p;
    }

    public static PytorchClassifier create(Context context, String modelPath, int inputsize) throws IOException {
        PytorchClassifier p = new PytorchClassifier();
        p.module = Module.load(assetFilePath(context, modelPath));
        p.inputsize = inputsize;
        return p;
    }

    @Override
    public List<Recognition> recognizeImage(float[] pixels) {

        // preparing input tensor
        final Tensor inputTensor = Tensor.fromBlob(pixels, new long[] {1, 1, inputsize, inputsize});

        // running the model
        final Tensor outputTensor = module.forward(IValue.from(inputTensor)).toTensor();

        // getting tensor content as java array of floats
        final float[] scores = outputTensor.getDataAsFloatArray();

        // searching for the index with maximum score
        float maxScore = -Float.MAX_VALUE;
        int maxScoreIdx = -1;
        for (int i = 0; i < scores.length; i++) {
            if (scores[i] > maxScore) {
                maxScore = scores[i];
                maxScoreIdx = i;
            }
        }
        ArrayList<Recognition> recognitions = new ArrayList<>();
        recognitions.add(new Recognition(maxScoreIdx + "", maxScoreIdx + "", maxScore, null));
        return recognitions;
    }

    @Override
    public void enableStatLogging(boolean debug) {
        isLog = debug;
    }

    @Override
    public String getStatString() {
        return null;
    }

    @Override
    public void close() {
        if (module != null) {
            module.destroy();
        }
        module = null;
    }

    public static String assetFilePath(Context context, String assetName) throws IOException {
        File file = new File(context.getFilesDir(), assetName);
        if (file.exists() && file.length() > 0) {
            return file.getAbsolutePath();
        }

        return saveStreamAsFile(context, context.getAssets().open(assetName), file);
    }

    public static String saveStreamAsFile(Context context, InputStream in, File dest) throws IOException {
        try (InputStream is = in) {
            try (OutputStream os = new FileOutputStream(dest)) {
                byte[] buffer = new byte[4 * 1024];
                int read;
                while ((read = is.read(buffer)) != -1) {
                    os.write(buffer, 0, read);
                }
                os.flush();
            }
            return dest.getAbsolutePath();
        }
    }
}
